package com.zshan.clinic.web.filter;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.lang3.StringUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class CustomHttpServletRequest extends HttpServletRequestWrapper {

    private final Map<String, String> headers = new HashMap<>();

    private Map<String, Object> customParams = new HashMap<>();

    // 缓存请求体字节数组
    private byte[] cachedBody;


    public CustomHttpServletRequest(HttpServletRequest request) throws IOException {
        super(request);
        cacheInputStream(request);
        // 初始化参数
        this.customParams = getRequestParams(this);
    }

    // 缓存请求体
    private void cacheInputStream(HttpServletRequest request) throws IOException {
        InputStream requestInputStream = request.getInputStream();
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        byte[] data = new byte[1024];
        int nRead;
        while ((nRead = requestInputStream.read(data, 0, data.length)) != -1) {
            buffer.write(data, 0, nRead);
        }
        buffer.flush();
        cachedBody = buffer.toByteArray();
    }

    @Override
    public ServletInputStream getInputStream() {
        ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(this.cachedBody);
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return byteArrayInputStream.available() == 0;
            }
            @Override
            public boolean isReady() {
                return true;
            }
            @Override
            public void setReadListener(ReadListener listener) {
                // not implemented
            }
            @Override
            public int read() {
                return byteArrayInputStream.read();
            }
        };
    }

    @Override
    public BufferedReader getReader() {
        return new BufferedReader(new InputStreamReader(this.getInputStream(), StandardCharsets.UTF_8));
    }

    // 你的其他代码保持不变
    public void addHeader(String name, String value) {
        headers.put(name, value);
    }

    public void setParams(Map<String, Object> params) {
        if (params != null) {
            this.customParams = params;
            // 如果是 application/json 且是 POST，请更新 cachedBody
            if ("POST".equalsIgnoreCase(this.getMethod()) &&
                    getContentType() != null &&
                    getContentType().toLowerCase().contains("application/json")) {
                try {
                    ObjectMapper mapper = new ObjectMapper();
                    String json = mapper.writeValueAsString(this.customParams);
                    this.cachedBody = json.getBytes(StandardCharsets.UTF_8);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }


    @Override
    public String getHeader(String name) {
        if (headers.containsKey(name)) {
            return headers.get(name);
        }
        return super.getHeader(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        List<String> names = Collections.list(super.getHeaderNames());
        names.addAll(headers.keySet());
        return Collections.enumeration(new HashSet<>(names)); // 去重
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        List<String> list = Collections.list(super.getHeaders(name));
        if (headers.containsKey(name)) {
            list.add(headers.get(name));
        }
        return Collections.enumeration(list);
    }

    @Override
    public String getParameter(String name) {
        Object value = customParams.get(name);
        if (value instanceof String) {
            return (String) value;
        } else if (value instanceof List) {
            List<?> list = (List<?>) value;
            return list.isEmpty() ? null : String.valueOf(list.get(0));
        } else if (value != null) {
            return String.valueOf(value);
        }
        return super.getParameter(name);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> paramMap = new HashMap<>();
        for (Map.Entry<String, Object> entry : customParams.entrySet()) {
            Object value = entry.getValue();
            if (value instanceof String) {
                paramMap.put(entry.getKey(), new String[]{(String) value});
            } else if (value instanceof List) {
                List<?> list = (List<?>) value;
                String[] arr = list.stream().map(String::valueOf).toArray(String[]::new);
                paramMap.put(entry.getKey(), arr);
            } else if (value != null) {
                paramMap.put(entry.getKey(), new String[]{String.valueOf(value)});
            }
        }
        return paramMap;
    }

    @Override
    public String[] getParameterValues(String name) {
        Object value = customParams.get(name);
        if (value instanceof String) {
            return new String[]{(String) value};
        } else if (value instanceof List) {
            List<?> list = (List<?>) value;
            return list.stream().map(String::valueOf).toArray(String[]::new);
        } else if (value != null) {
            return new String[]{String.valueOf(value)};
        }
        return super.getParameterValues(name);
    }

    @Override
    public Enumeration<String> getParameterNames() {
        return Collections.enumeration(customParams.keySet());
    }

    /**
     * 获取请求参数，兼容 GET 表单 和 POST JSON
     */
    public Map<String, Object> getRequestParams(HttpServletRequest request) {
        Map<String, Object> params = new HashMap<>();

        // GET / 表单 POST 参数
        Map<String, String[]> parameterMap = request.getParameterMap();
        for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
            String key = entry.getKey();
            String[] values = entry.getValue();
            if (values.length == 1) {
                params.put(key, values[0]);
            } else {
                params.put(key, Arrays.asList(values));
            }
        }

        // JSON 类型的 POST 参数，改为从缓存中读取，避免重复读取流
        if ("POST".equalsIgnoreCase(request.getMethod()) &&
                request.getContentType() != null &&
                request.getContentType().contains("application/json")) {
            try {
                // 用缓存的body转换字符串，而不是调用 request.getReader()
                String body = new String(this.cachedBody, StandardCharsets.UTF_8);
                if (StringUtils.isNotBlank(body)) {
                    ObjectMapper mapper = new ObjectMapper();
                    Map<String, Object> jsonMap = mapper.readValue(body, Map.class);
                    params.putAll(jsonMap);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return params;
    }
}
